package week07.springdsrouter;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.transaction.support.TransactionSynchronizationManager;

import javax.sql.DataSource;
import java.util.List;
import java.util.concurrent.ThreadLocalRandom;

/**
 * @author edward
 * @since 2021-08-07
 */
public class ReadWriteDataSourceRouter extends AbstractRoutingDataSource {
    @Override
    protected Object determineCurrentLookupKey() {
        if (TransactionSynchronizationManager.isCurrentTransactionReadOnly()) {
            System.out.println("Datasource from slaves");
            return DataSourceRoutingApplicationConfig.SLAVES;
        }
        System.out.println("Datasource from master");
        return DataSourceRoutingApplicationConfig.MASTER;
    }

    @Override
    protected DataSource resolveSpecifiedDataSource(Object dataSource) throws IllegalArgumentException {
        if (dataSource instanceof List) {
            //noinspection unchecked
            List<DataSource> dataSourceList = (List<DataSource>) dataSource;
            // 随机获取一个 slave 服务
            return dataSourceList.get(ThreadLocalRandom.current().nextInt(dataSourceList.size()));
        }
        return super.resolveSpecifiedDataSource(dataSource);
    }
}
